# -*- coding: utf-8 -*-  
'''
实在么得办法，找不到与知识库对应的数据集，自己造一个。。。

Created on 2021年10月1日
@author: luoyi
'''
import random
import json
import tqdm

import utils.conf as conf


#    根据关系生成疑问句
class QuestionsGenerator:
    def __init__(self, ds_baidu_iter, question_schemas_path=conf.DATASET_BAIDU.get_question_schemas_path(), n_shuffle=1024):
        '''
            @param question_schemas_path: 提问规则文件
            @param shuffle_buffer: 先随机写入n_shuffle个文件中，在按顺序合并到一个文件中
        '''
        self._ds_baidu_iter = ds_baidu_iter
        self._n_shuffle = n_shuffle
        self._qi_dict = self.parse_question_schemas(question_schemas_path)
        pass
    
    #    解析提问规则文件
    def parse_question_schemas(self, question_schemas_path):
        qi_dist = {}
        with open(file=question_schemas_path, mode='r', encoding='utf-8') as fr:
            line = fr.readline()
            
            while (line is not None):
                line = line.strip().replace('\n', '')
                if (len(line) == 0): break
                
                j = json.loads(line)
                
                qi = QuetionItem(j)
                qi_dist[qi._predicate] = qi
                
                line = fr.readline()
                pass
            pass
        
        return qi_dist

    
    #    生成疑问句
    def create_questions(self, fpath=conf.DATASET_BAIDU.get_question_train_data_path(), log_total=173040):
        #    过程中用到的随机对象
        rdm = random.Random(self._n_shuffle)
        
        #    建目录
        conf.mkfiledir_ifnot_exises(fpath)
        shuffle_dir = fpath + "_shuffle/"
        conf.mkdir_ifnot_exises(shuffle_dir)
        
        #    所有待写的文件
        fws = [open(shuffle_dir + '/' + str(i), mode='w', encoding='utf-8') for i in range(self._n_shuffle)]
        
        #    遍历每个样本，并将生成的问句随机写入文件中
        for sample in tqdm.tqdm(iterable=self._ds_baidu_iter, desc='遍历baidu数据集.写入临时文件', total=log_total, leave=True, ncols=100):
            qs = self.question(sample)
            if qs is None: continue
            
            for q in qs:
                j = json.dumps(q)
            
                #    随机写入文件中
                fw = rdm.choice(fws)
                fw.write(j + "\n")
                pass
            pass
        #    关闭所有文件
        for fw in fws: fw.close()
        
        total = 0
        #    顺序合并所有文件内容
        fw = open(fpath, mode='w', encoding='utf-8')
        for i in tqdm.tqdm(iterable=range(self._n_shuffle), desc='write_tfrecord', total=self._n_shuffle, leave=True, ncols=100):
            with open(shuffle_dir + '/' + str(i), mode='r', encoding='utf-8') as fr:
                line = fr.readline()
                while (line is not None):
                    line = line.strip().replace('\n', '')
                    if len(line) == 0: break
                    
                    line = line.encode('utf-8')
                    line = line.decode('unicode-escape')
                    fw.write(line + '\n')
                    total += 1
                    line = fr.readline()
                    pass
                pass
            pass
        fw.close()
        
        #    删除临时文件
#         os.removedirs(shuffle_dir)
        
        return total
    
    #    生成问句
    def question(self, sample):
        '''格式，每行一个json:
            {
                sen: 问句, 
                predicate:提问的关系, 
                subject:{
                    word: 实体值。若q=True，则这里视为None。但一般会给真实值
                    loc: (在问句中首尾指针)。若q=True，则该属性为空
                    type: 实体类型
                    q: True|False，是否针对该实体进行提问
                },
                object:{
                    word: 实体值。若q=True，则这里视为None。但一般会给真实值
                    loc: (在问句中首尾指针)。若q=True，则该属性为空
                    type: 实体类型
                    q: True|False，是否针对该实体进行提问
                },
            }
        '''
        qs = []
        
        #    根据关系进行提问
        for spo in sample.spo_list:
            predicate = spo['predicate']
            qi = self._qi_dict.get(predicate)
            if (qi is None): continue
            
            #    生成针对subject和object的疑问句
            q_subject, q_object = qi.create_question(v_subject=spo['subject']['word'], v_object=spo['object']['word'])
            
            qs = qs + q_subject
            qs = qs + q_object
            pass
        
        return qs
    pass


#    遍历问题
class QuestionIterator:
    def __init__(self):
        pass
    
    #    遍历迭代器
    def data_iterator(self, fpath=conf.DATASET_BAIDU.get_tfrecord_train_data_path(), count=-1):
        '''格式，每行一个json:
            {
                sen: 问句, 
                predicate:提问的关系, 
                subject:{
                    word: 实体值。若q=True，则这里视为None。但一般会给真实值
                    loc: (在问句中首尾指针)。若q=True，则该属性为空
                    type: 实体类型
                    q: True|False，是否针对该实体进行提问
                },
                object:{
                    word: 实体值。若q=True，则这里视为None。但一般会给真实值
                    loc: (在问句中首尾指针)。若q=True，则该属性为空
                    type: 实体类型
                    q: True|False，是否针对该实体进行提问
                },
            }
        '''
        i = 0
        with open(file=fpath, mode='r', encoding='utf-8') as fr:
            line = fr.readline()
            
            while (line is not None):
                line = line.strip().replace('\n', '')
                if (len(line) == 0): break
                
                try :
                    j = json.loads(line)
                except json.decoder.JSONDecodeError:
                    line = line.replace('\\', '|')
                    j = json.loads(line)
                    pass
                yield self.parse_json(j)
                
                if (count > 0 and i >= count): break
                
                line = fr.readline()
                i += 1
                pass
            pass
        pass
    
    #    解析json
    def parse_json(self, j):
        '''
            @return: sen, predicate, subject_(word, loc(可为空), type, q), object_
        '''
        sen = j.get('sen')
        predicate = j.get('predicate')
        ks = j['subject']
        subject_ = (ks.get('word'), ks.get('loc'), ks.get('type'), ks.get('q'))
        ko = j['object']
        object_ = (ko.get('word'), ko.get('loc'), ko.get('type'), ko.get('q'))
        return sen, predicate, subject_, object_
    
    pass


#    生成疑问句工具
class QuetionItem:
    def __init__(self, relu_text):
        self.parse_relu(relu_text)
        pass
    
    #    解析规则文本
    def parse_relu(self, relu_text):
        #    关系，subject类型，object类型
        self._predicate = relu_text['predicate']
        self._subject_type = relu_text['subject_type']
        self._object_type = relu_text['object_type']
        
        #    疑问句规则
        self._subject_ask_object = relu_text.get('subject_ask_object')
        self._object_ask_subject = relu_text.get('object_ask_subject')
        pass

    #    生成疑问句
    def create_question(self, v_subject, v_object):
        #    解析向object的提问
        q_sao = []
        if (self._subject_ask_object is not None):
            for sao in self._subject_ask_object:
                try: 
                    loc_s = sao.index('[subject]')
                    loc_e = loc_s + len(v_subject)
                    loc = (loc_s, loc_e)
                    
                    subject_ = {'word': v_subject, 'loc':loc, 'type':self._subject_type, 'q':False}
                    object_ = {'word': v_object, 'type':self._object_type, 'q':True}
                    sen = sao.replace('[subject]', v_subject)
                    
                    q = {'predicate':self._predicate, 'sen':sen, 'subject':subject_, 'object':object_}
                    q_sao.append(q)
                    
                except ValueError: 
                    continue
                pass
            pass
        
        #    解析向subject的提问
        q_oas = []
        if (self._object_ask_subject is not None):
            for oas in self._object_ask_subject:
                try: 
                    loc_s = oas.index('[object]')
                    loc_e = loc_s + len(v_object)
                    loc = (loc_s, loc_e)
                    
                    subject_ = {'word': v_subject, 'type':self._subject_type, 'q':True}
                    object_ = {'word': v_object, 'loc':loc, 'type':self._object_type, 'q':False}
                    sen = oas.replace('[object]', v_object)
                    
                    q = {'predicate':self._predicate, 'sen':sen, 'subject':subject_, 'object':object_}
                    q_sao.append(q)
                    
                except ValueError: 
                    continue
                pass        
            pass
        return q_sao, q_oas
    pass

